import pygame
import real_env
import flask
from flask import Flask, send_file
import io
from PIL import Image
import time

app = Flask(__name__)

pygame.init()
ENV = real_env.GuidewireEnv('real_x_ray')
ENV.max_steps = 50


ACTIONS = {
    "push": 0,
    "pull": 1,
    "bend_cw": 2,
    "bend_ccw": 3,
    "stop": 4
}
ACT = ACTIONS["stop"]

LAST_ACTION_TIME = time.time()

@app.route('/capture', methods=['GET'])
def capture():
    pil_image = Image.frombytes('RGB', ENV.display.get_size(), pygame.image.tostring(ENV.display, 'RGB'))
    # 将PIL图像保存到字节流中
    img_io = io.BytesIO()
    pil_image.save(img_io, 'PNG')
    img_io.seek(0)
    return send_file(img_io, mimetype='image/png')

@app.route('/reset', methods=['POST'])
def reset():
    global ENV
    ENV.reset()
    return {"status": "reset"}

@app.route('/action', methods=['POST'])
def action():
    global ACT, ENV, LAST_ACTION_TIME
    time_diff = 0
    current_time = time.time()
    
    time_diff = current_time - LAST_ACTION_TIME
    LAST_ACTION_TIME = current_time
    if time_diff < 1.:
        return {"status": "busy", "past_time": f"{time_diff:.1f}" }
    
    ACT = ACTIONS[flask.request.form['action']]
    ENV.step(ACT)
    return {"status": "success", "past_time": f"{time_diff:.1f}" }


@app.route('/action_fast', methods=['POST'])
def action_fast():
    global ACT, ENV, LAST_ACTION_TIME
    ACT = ACTIONS[flask.request.form['action']]
    ENV.step(ACT)
    return {"status": "success" }


if __name__ == '__main__':
    app.run(host='0.0.0.0', port=5003)